import os
import torch
from torchvision import transforms
from PIL import Image
from torch.utils.data import DataLoader, Dataset
from dataset.Data_Processing import get_total_list, predicate_to_seq, data_processing


class LOAD_DATA(Dataset):
    def __init__(self, img_path, file_path, mode='train'):
        # print(img_path, file_path)
        self.img_path = img_path
        self.file_path = file_path
        self.mode = mode
        self.img_list = os.listdir(self.img_path)
        self.features_seq_list, self.name_list, self.class_list = data_processing(file_path)

    def __getitem__(self, index):
        img_name = self.name_list[index]
        img_path = os.path.join(self.img_path, self.name_list[index])
        img = Image.open(img_path)
        try:
            pos = self.name_list.index(img_name)
            features_seq = self.features_seq_list[pos]
            label = self.class_list[pos]
        except Exception as e:
            print("error", e)
        if self.mode == 'train':
            transform = transforms.Compose([
                transforms.Resize((224, 224)),
                transforms.RandomHorizontalFlip(p=0.5),
                transforms.ToTensor()
            ])
        else:
            transform = transforms.Compose([
                transforms.Resize((224, 224)),
                transforms.ToTensor()
            ])
        img = transform(img)
        return img, torch.tensor(features_seq), label, img_name

    def __len__(self):
        return len(self.name_list)


# if __name__ == '__main__':
#     train_dataset = LOAD_DATA("../data/images/imgs", "../data/bladder_train.json", mode='imgs')
#     train_loader = DataLoader(train_dataset, batch_size=12, shuffle=True, num_workers=0)
#     print(next(iter(train_loader)))
